﻿#pragma once

#include "InitialCondition.h"

struct IsoVortexIC : WorkletInitCondition<IsoVortexIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    ValueType value;
    FloatDefault t = 0;
    Vec3f _vel{ 1, 1, 0 };
    FloatDefault x = p[0] - _vel[0] * t;
    FloatDefault y = p[1] - _vel[1] * t;
    FloatDefault z = p[2] - _vel[2] * t;
    FloatDefault gam = 1.4, gamm1 = gam - 1, epi = 5.;
    FloatDefault xb, yb, r2;

    FloatDefault PI = vtkm::Pi();
    xb = x + 5;
    yb = y + 5;
    r2 = xb * xb + yb * yb;

    FloatDefault u = _vel[0] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * (-yb);
    FloatDefault v = _vel[1] + epi / (2. * PI) * exp(0.5 * (1. - r2)) * xb;
    FloatDefault T = 1. - gamm1 * epi * epi / (8 * gam * PI * PI) * exp(1. - r2);
    FloatDefault density = pow(T, 1 / gamm1);
    FloatDefault pre = pow(density, gam);

    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }
};


struct CoutteFlowIC : WorkletInitCondition<CoutteFlowIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1, mach = 0.2;
    FloatDefault u = 1;
    FloatDefault v = 0;
    FloatDefault density = 1;
    FloatDefault pre = 1. / (gam * mach * mach);

    ValueType value;
    value[0] = density;
    value[1] = density * u;
    value[2] = density * v;
    value[3] = density * 0;
    value[4] = pre / gamm1 + 0.5 * density * (u * u + v * v);

    return value;
  }
};


struct Riemann2DIC : WorkletInitCondition<Riemann2DIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    ValueType primary_value;
    auto x = p[0];
    auto y = p[1];
    if (x > 0.5 && y > 0.5)
    {
      primary_value = ValueType{ 1.5, 0, 0, 0, 1.5 };
    }
    else if (x < 0.5 && y > 0.5)
    {
      primary_value = ValueType{ 0.5323, 1.206, 0, 0, 0.3 };
    }
    else if (x < 0.5 && y < 0.5)
    {
      primary_value = ValueType{ 0.138, 1.206, 1.206, 0, 0.029 };
    }
    else // if (x > 0.5 && y < 0.5)
    {
      primary_value = ValueType{ 0.5323, 0, 1.206, 0, 0.3 };
    }

    return PrimaryToConservation(primary_value);
  }

  VTKM_EXEC ValueType PrimaryToConservation(const ValueType& primary_value) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1;
    ValueType value;
    auto r = primary_value[0];
    auto u = primary_value[1];
    auto v = primary_value[2];
    auto w = primary_value[3];
    auto p = primary_value[4];

    value[0] = r;
    value[1] = r * u;
    value[2] = r * v;
    value[3] = r * w;
    value[4] = p / gamm1 + 0.5 * r * (u * u + v * v);
    return value;
  }
};

class IsoVortexInitialCondition : public InitialCondition
{
public:
  IsoVortexInitialCondition(){};
  ~IsoVortexInitialCondition(){};

  virtual void Execute() override {}

private:
};